from data_output import write_to_py_file
from freq_gen import freq_lis
import numpy as np
import os

#频谱数据文件名
default_freq_file_name = 'harps_freq_res.py'
#乐谱数据文件名
default_time_file_name = 'yinjie_3_viola.py'

def main(freq_file_name=default_freq_file_name, time_file_name=default_time_file_name, auto_instrument_type=0, output_py_name=''):
    # 读取文件内容
    with open(freq_file_name, "r", encoding="utf-8") as freq_filedata:
        freq_fileinfo = freq_filedata.read()

    with open(time_file_name, "r", encoding="utf-8") as time_filedata:
        time_fileinfo = time_filedata.read()

    # 创建一个空字典作为全局命名空间
    global_namespace_1 = {}
    exec(freq_fileinfo, global_namespace_1)

    global_namespace_2 = {}
    exec(time_fileinfo, global_namespace_2)

    # 假设数组变量名为 data
    freq_res = global_namespace_1.get("res_list")
    time_res = global_namespace_2.get("res_list")

    instrument_type_code = list(np.zeros(8))
    if auto_instrument_type == 0:

        print('instrument_type?')
        instrument_type = input('>')

        if instrument_type == '提琴' or instrument_type == '1':
            instrument_type_code[1] = 1
        elif instrument_type == '钢琴' or instrument_type == '2':
            instrument_type_code[2] = 1
        elif instrument_type == '吉他' or instrument_type == '3':
            instrument_type_code[3] = 1
        elif instrument_type == '笛子' or instrument_type == '4':
            instrument_type_code[4] = 1
        elif instrument_type == '铜管' or instrument_type == '5':
            instrument_type_code[5] = 1
        elif instrument_type == 'misc' or instrument_type == '0':
            instrument_type_code[0] = 1
        else:
            instrument_type_code[0] = 1

    elif auto_instrument_type == 1:
        if 'violin' in freq_file_name:
            instrument_type_code[1] = 1
        elif 'piano' in freq_file_name:
            instrument_type_code[2] = 1
        elif 'guitar' in freq_file_name:
            instrument_type_code[3] = 1
        elif 'flute' in freq_file_name:
            instrument_type_code[4] = 1
        elif 'brass' in freq_file_name:
            instrument_type_code[5] = 1
        elif 'misc' in freq_file_name:
            instrument_type_code[0] = 1
        else:
            instrument_type_code[0] = 1

    data_set = []

    for index, item in enumerate(freq_res):
        if index+1 < len(freq_res):
            sample_time_cross = freq_res[index+1][-1] - freq_res[index][-1]
        for time_win in time_res:
            time_cross = time_win[1] - time_win[0]
            if item[-1]*1000 > time_win [0] + time_cross * 0.05 and (item[-1]+sample_time_cross)*1000 < time_win[1] - time_cross * 0.05:
                data_set.append([item[:-1], item[-1], time_win[2], instrument_type_code])

    for item in data_set:
        dect_freq_lis = item[2].split('+')
        freq_data = np.zeros(88)
        for freq in dect_freq_lis:
            for n in range(len(freq_lis)):
                if freq == freq_lis[n][2]:
                    freq_data[n] = 1
        item[2] = list(freq_data)


    if output_py_name == '':
        output_py_name, _ = os.path.splitext(freq_file_name)
    output_py_name = 'dataset_from_' + output_py_name + '.py'
    write_to_py_file(output_py_name=output_py_name, output_data=data_set, write_mode='str')

if __name__ == "__main__":
    main()